from llmtcg.store.codebase.base import CodeBase, FileChunk, FileChunks, Pos, PosRange
from llmtcg.utils.treesitter import TreeUtils, VisitHint

from loguru import logger

from dataclasses import dataclass, field
from os import listdir, path, walk
from typing import Dict, List, Set, Optional
from time import time
from itertools import chain

import tree_sitter_c
import tree_sitter


clanguage = tree_sitter.Language(tree_sitter_c.language())
cparser = tree_sitter.Parser()
cparser.language = clanguage


class CTreeHelper:

    @staticmethod
    def range_of(node: tree_sitter.Node) -> PosRange:
        return PosRange(
            Pos(*(node.start_point)),
            Pos(*(node.end_point)),
        )

    @staticmethod
    def declarator_name(declarator: tree_sitter.Node) -> Optional[str]:
        declarator_text = TreeUtils.textof(declarator)
        if declarator.type.startswith('abstract_') and declarator.type.endswith('declarator'):
            # including:
            # abstract_function_declarator, abstract_pointer_declarator, abstract_array_declarator, ...
            return None
        elif declarator.type in ('pointer_declarator', 'array_declarator', 'function_declarator'):
            sub_declarator = declarator.child_by_field_name('declarator')
            if sub_declarator is None:
                return None

            return CTreeHelper.declarator_name(sub_declarator)
        elif declarator.type in ('identifier', 'type_identifier', 'field_identifier'):
            return declarator_text
        elif declarator.type == 'primitive_type':
            return declarator_text
        elif declarator.type == 'parenthesized_declarator':
            return CTreeHelper.declarator_name(declarator.children[1])
        elif declarator.type == 'struct_specifier':
            name_node = declarator.child_by_field_name('name')
            if name_node is None:
                return None
            else:
                struct_name = CTreeHelper.declarator_name(name_node)
                if struct_name is None:
                    return None
                else:
                    return f"struct#{struct_name}"
        elif declarator.type == 'type_definition':
            return CTreeHelper.declarator_name(TreeUtils.get_child(declarator, 'declarator'))
        elif declarator.type == 'ERROR':
            return None

        raise NotImplementedError(declarator, declarator_text)

    @staticmethod
    def collect_type_names(ty: tree_sitter.Node) -> List[str]:
        typenames: List[str] = []

        def __visit(node: tree_sitter.Node) -> Optional[VisitHint]:
            nonlocal typenames
            if node.type == 'struct_specifier':
                name = node.child_by_field_name('name')
                if name is not None:
                    typenames.append(f"struct#{TreeUtils.textof(name)}")
            elif node.type == 'primitive_type':
                typenames.append(TreeUtils.textof(node))

            return None

        TreeUtils.visit(ty, __visit)
        return typenames


@dataclass
class CDeclaration:
    name: Optional[str]
    range: PosRange

    @classmethod
    def from_tree_node(self, declarator: tree_sitter.Node) -> 'CDeclaration':

        ROOT_TREE_TYPES = set([
            'declaration',
            'function_definition', 'type_definition',
            'struct_specifier'
        ])
        decl_node = declarator
        while decl_node.type not in ROOT_TREE_TYPES and decl_node.parent is not None:
            decl_node = decl_node.parent

        start_pos = Pos(*decl_node.start_point)
        end_pos = Pos(*decl_node.end_point)

        name = CTreeHelper.declarator_name(declarator)
        return CDeclaration(name, PosRange(start_pos, end_pos))

    def set_file_summary(self, fs: 'CFileSummary') -> None:
        setattr(self, "_file_summary", fs)

    def get_file_summary(self) -> 'CFileSummary':
        if not hasattr(self, "_file_summary"):
            raise TypeError(f"declaration is not properly typed")

        fs = getattr(self, "_file_summary")
        assert isinstance(fs, CFileSummary)

        return fs

    def to_file_chunk(self) -> FileChunk:
        fc = FileChunk(self.get_file_summary().filename, self.range.start, self.range.end)
        return fc

    def __hash__(self) -> int:
        return id(self)


@dataclass
class CMacro(CDeclaration):

    @classmethod
    def from_tree_node(self, macro: tree_sitter.Node) -> 'CMacro':
        macro_node = TreeUtils.get_child(macro, 'name')
        macro_name = TreeUtils.textof(macro_node)
        return CMacro(macro_name, CTreeHelper.range_of(macro))

    def __hash__(self) -> int:
        return id(self)


@dataclass
class CMacroFunction(CDeclaration):

    @classmethod
    def from_tree_node(self, macro: tree_sitter.Node) -> 'CMacroFunction':
        macro_node = TreeUtils.get_child(macro, 'name')
        macro_name = TreeUtils.textof(macro_node)
        return CMacroFunction(macro_name, CTreeHelper.range_of(macro))

    def __hash__(self) -> int:
        return id(self)


@dataclass
class CTypedef(CDeclaration):

    @classmethod
    def from_tree_node(self, declarator: tree_sitter.Node) -> 'CTypedef':
        parent_decl = CDeclaration.from_tree_node(declarator)
        return CTypedef(parent_decl.name, parent_decl.range)

    def __hash__(self) -> int:
        return id(self)


@dataclass(repr=False)
class CParameter(CDeclaration):

    related_type_names: List[str]

    @classmethod
    def from_tree_node(self, node: tree_sitter.Node) -> 'CParameter':
        assert node.type == 'parameter_declaration'
        param_name: Optional[str] = None
        declarator = node.child_by_field_name('declarator')
        if declarator is not None:
            param_name = CTreeHelper.declarator_name(declarator)

        param_range = CTreeHelper.range_of(node)

        type_names = list()
        type_node = node.child_by_field_name('type')
        if type_node is not None:
            type_names = CTreeHelper.collect_type_names(type_node)

        param = CParameter(param_name, param_range, type_names)
        return param

    def __hash__(self) -> int:
        return id(self)


@dataclass
class CFunction(CDeclaration):

    parameters: List[CParameter]
    defined: bool
    callee_names: Set[str] = field(default_factory=set)

    @classmethod
    def from_tree_node(cls, node: tree_sitter.Node) -> 'CFunction':

        body_node: Optional[tree_sitter.Node] = None
        if node.type == 'function_definition':
            declarator = TreeUtils.get_child(node, "declarator")
            while declarator.type == 'pointer_declarator':
                # a function definition cannot be defined by a pointer declarator, e.g.
                # `void (*f) { ... }` never happens. but in treesitter we find a case,
                # if the function's return type is a struct pointer, it happens to be
                # treated as a pointer declarator here
                declarator = TreeUtils.get_child(declarator, 'declarator')

            body_node = TreeUtils.get_child(node, "body")
            defined = True
        elif node.type == 'function_declarator':
            declarator = node
            defined = False
        else:
            raise NotImplementedError(node.sexp())

        parnet_decl = CDeclaration.from_tree_node(declarator)
        parameters: List[CParameter] = []
        params_node = TreeUtils.get_child(declarator, 'parameters')
        for c in params_node.children:
            if c.type == 'parameter_declaration':
                parameters.append(CParameter.from_tree_node(c))

        fdecl = CFunction(
            parnet_decl.name, parnet_decl.range,
            parameters,
            defined
        )

        def visit_func_body(node: tree_sitter.Node) -> Optional[VisitHint]:
            if node.type == 'call_expression':
                callee = TreeUtils.get_child(node, "function")
                if callee.type == 'identifier':
                    fdecl.callee_names.add(TreeUtils.textof(callee))

            return None

        if body_node is not None:
            TreeUtils.visit(body_node, visit_func_body)

        return fdecl

    def parameter_type_names(self) -> List[str]:
        return list(chain(*(p.related_type_names for p in self.parameters)))

    def __hash__(self) -> int:
        return id(self)


@dataclass
class CFileChunksBuilder:
    fcs: FileChunks
    added_items: Set[CDeclaration] = field(default_factory=set, init=False)

    def add_type_decl(self, decl: CDeclaration) -> None:
        if decl in self.added_items:
            return None

        self.fcs.add(decl.to_file_chunk())

    def add_function(self, function: CFunction) -> None:
        if function in self.added_items:
            return

        self.added_items.add(function)
        self.fcs.add(function.to_file_chunk())
        for type_name in function.parameter_type_names():
            type_decls = function.get_file_summary().find_type_decls(type_name)
            for ty_decl in type_decls:
                self.add_type_decl(ty_decl)

    def build(self) -> FileChunks:
        return self.fcs


@dataclass
class CFileSummary:
    filename: str
    includes: List[str] = field(default_factory=list, init=False)
    global_decls: Dict[str, List[CDeclaration]] = field(default_factory=dict, init=False, repr=False)

    def set_codebase(self, codebase: 'CCodeBase') -> None:
        setattr(self, "_codebase", codebase)

    def get_codebase(self) -> 'CCodeBase':
        cb = getattr(self, "_codebase")
        if not isinstance(cb, CCodeBase):
            raise TypeError("cfilesummary is not properly typed")

        return cb

    @classmethod
    def load_file(cls, filename: str) -> 'CFileSummary':
        with open(filename, 'rb') as f:
            content = f.read()

        return cls.load_string(content, filename)

    @classmethod
    def load_string(cls, content: bytes, filename: str = "_") -> 'CFileSummary':
        tree = cparser.parse(content)
        cfs = CFileSummary(filename)
        cfs.feed_tree(tree, filename)
        return cfs

    def get_included_file_summaries(self, recursive: bool = False) -> List['CFileSummary']:
        file_summaries = []
        codebase = self.get_codebase()
        for header in self.includes:
            file_summaries.extend(codebase.find_file_summaries(header))

        return file_summaries

    def find_type_decls(self, type_name: str, with_includes: bool = True) -> List[CDeclaration]:
        file_summaries: Set[CFileSummary] = set([self])
        if with_includes:
            file_summaries.update(self.get_included_file_summaries(recursive=True))

        decls: List[CDeclaration] = []
        for fs in file_summaries:
            for decl in fs.global_decls.get(type_name, []):
                if isinstance(decl, (CTypedef, CMacro)):
                    decls.append(decl)

        return decls

    def __hash__(self) -> int:
        return hash((type(self), self.filename))

    def feed_tree(self, tree: tree_sitter.Tree, filename: str) -> None:
        """ Build symbol map of this file """

        def register_decl(decl: CDeclaration) -> None:
            name = decl.name
            if name is None:
                return

            if name not in self.global_decls:
                self.global_decls[name] = []

            decl.set_file_summary(self)
            self.global_decls[name].append(decl)

        def register_function(fnode: tree_sitter.Node) -> None:
            try:
                fdecl = CFunction.from_tree_node(fnode)
                register_decl(fdecl)
            except Exception as ex:
                logger.warning(f"failed parsing function at line {fnode.start_point[0]}, {filename}: {ex}")

        def register_type_definition(tydefnode: tree_sitter.Node) -> None:
            declarator = TreeUtils.get_child(tydefnode, 'declarator')
            if declarator.type == 'primitive_type':
                tydecl = CTypedef.from_tree_node(tydefnode)
                register_decl(tydecl)

        def __visit(n: tree_sitter.Node) -> Optional[VisitHint]:
            """ Recursively visit the ast to find global symbols """
            if n.type in ("function_definition", "function_declarator"):
                register_function(n)
                return VisitHint.STEP_OVER
            elif n.type == 'struct_specifier':
                if n.child_by_field_name('body') is not None:
                    register_decl(CTypedef.from_tree_node(n))

                return VisitHint.STEP_OVER
            elif n.type == 'type_definition':
                register_type_definition(n)
                return VisitHint.STEP_OVER
            elif n.type == "preproc_include":
                path_node = n.child_by_field_name("path")
                if path_node is not None:
                    # the original text in this node is a string liter "..."
                    self.includes.append(TreeUtils.textof(path_node)[1:-1])
            elif n.type == 'preproc_def':
                register_decl(CMacro.from_tree_node(n))
            elif n.type == 'preproc_function_def':
                register_decl(CMacroFunction.from_tree_node(n))

            return None

        TreeUtils.visit(tree.root_node, __visit)


class CCodeBase(CodeBase):

    def __init__(self) -> None:
        super().__init__()

        self.files: Dict[str, CFileSummary] = dict()

    def clear(self) -> None:
        self.files.clear()

    def load_string(self, filename: str, content: bytes) -> None:
        fs = CFileSummary.load_string(content, filename)
        self.files[filename] = fs

    def load_directory(self, directory: str) -> None:
        start_time = time()
        for root, _, files in walk(directory, followlinks=True):
            for file in files:
                if file.endswith((".c", ".h")):
                    filename = path.join(root, file)
                    fs = CFileSummary.load_file(filename)
                    fs.set_codebase(self)
                    self.files[filename] = fs

        cost_time = time() - start_time
        logger.info(f"loading {directory} finished in {cost_time:.2f} seconds")

    ###########################################################################
    # Query functionality
    ###########################################################################

    def find_file_summaries(self, filename: str) -> List[CFileSummary]:
        """ The function search all the loaded files and find the ones that
        partially matches the filename as suffix. For example, `/path/to/linux/a.h`
        is a match of `linux/a.h`. """
        file_summaries = []
        for name, fs in self.files.items():
            if name.endswith(filename):
                file_summaries.append(fs)

        return file_summaries

    def find_functions(self, function_name: str) -> List[CFunction]:
        functions: List[CFunction] = []
        for file in self.files.values():
            for decl in file.global_decls.get(function_name, []):
                if isinstance(decl, CFunction):
                    functions.append(decl)

        return functions

    def find_declarations(self, decl_name: str) -> List[CDeclaration]:
        decls: List[CDeclaration] = []
        for file in self.files.values():
            decls.extend(file.global_decls.get(decl_name, []))

        return decls

    def find_chunks_by_function(self, function_name: str, path_pattern: Optional[str] = None) -> FileChunks:
        fcs = FileChunks()
        fcs_builder = CFileChunksBuilder(fcs)
        for func in self.find_functions(function_name):
            # todo: if the function's location do not match the path_pattern, ignore it
            fcs_builder.add_function(func)

        return fcs
